﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Windows.Forms;

namespace VersionComparer.UI.ExtendedTreeView
{
    public class MultiSelectTreeView : TreeView
    {
        #region fields
        SelectableTreeNode referenceShiftNode;
        SelectableTreeNode lastNode;
        List<SelectableTreeNode> selectedNodes = new List<SelectableTreeNode>();
        bool preventMultiSelect;
        #endregion fields

        #region constructors
        public MultiSelectTreeView()
            : base()
        {
            AutoSelectionEnabled = true;
        }
        #endregion constructors

        #region properties
        public bool AutoSelectionEnabled
        {
            get;
            set;
        }
        public TreeNode[] SelectedNodes
        {
            get
            {
                List<TreeNode> result = new List<TreeNode>();
                foreach (var node in this.selectedNodes)
                {
                    if (node.Selected)
                        result.Add(node);
                }
                if (SelectedNode != null && Focused)
                {
                    if (!result.Contains(SelectedNode))
                        result.Add(SelectedNode);
                }
                return result.ToArray();
            }
        }
        #endregion properties

        public void UnselectAll()
        {
            referenceShiftNode = null;
            if (selectedNodes.Count > 0)
            {
                BeginUpdate();
                List<TreeNode> nodesToUnselect = new List<TreeNode>(selectedNodes);
                foreach (SelectableTreeNode node in nodesToUnselect)
                {
                    if (node != null)
                        node.Selected = false;
                }
                selectedNodes.Clear();
                EndUpdate();
            }
        }
        public void SelectNode (TreeNode selectedNode)
        {
            lastNode = null;
            bool previousPreventMultiSelect = preventMultiSelect;
            preventMultiSelect = true;
            SelectedNode = selectedNode;
            preventMultiSelect = previousPreventMultiSelect;
        }
        public virtual void OnAfterMultiSelect(TreeViewEventArgs e)
        {
            SelectableTreeNode node = e.Node as SelectableTreeNode;
            if (node == null)
                return;
            if (node.Selected)
            {
                if (!selectedNodes.Contains(node))
                    selectedNodes.Add(node);
            }
            else
            {
                if (selectedNodes.Contains(node))
                    selectedNodes.Remove(node);
            }

        }
        private bool ModifierKeyIsPressed()
        {
            return (Control.ModifierKeys & Keys.Control) != 0 || (Control.ModifierKeys & Keys.Shift) != 0;
        }
        protected override void OnBeforeSelect(TreeViewCancelEventArgs e)
        {
            
            if (!preventMultiSelect && ModifierKeyIsPressed())
                BeginUpdate();
            lastNode = SelectedNode as SelectableTreeNode;
            if (!preventMultiSelect && lastNode != null && !ModifierKeyIsPressed ())
            {
                lastNode.SelectedEventEnabled = false;
                lastNode.Selected = false;
                lastNode.SelectedEventEnabled = true;
            }
            
            base.OnBeforeSelect(e);
            
        }
        protected override void OnAfterSelect(TreeViewEventArgs e)
        {
            try
            {
                bool unSelectAll = true;
                if ((Control.ModifierKeys & Keys.Control) > 0)
                {
                    unSelectAll = false;
                    if (Focused)
                    {
                        SelectableTreeNode node = e.Node as SelectableTreeNode;
                        if (node != null)
                        {
                            node.Selected = !node.Selected;
                            if (lastNode != null && !lastNode.Selected)
                                lastNode.Selected = true;
                        }
                    }
                    else
                        unSelectAll = true;
                }
                else
                {
                    if (!preventMultiSelect && (Control.ModifierKeys & Keys.Shift) > 0)
                    {
                        unSelectAll = false;
                        if (Focused)
                        {
                            if (AutoSelectionEnabled)
                            {
                                if (referenceShiftNode == null)
                                    referenceShiftNode = lastNode;
                                List<SelectableTreeNode> nodesToSelect = new List<SelectableTreeNode>(GetNodesBetweenNodes(e.Node as SelectableTreeNode, referenceShiftNode, true));
                                List<SelectableTreeNode> nodesToUnselect = new List<SelectableTreeNode>();
                                foreach (SelectableTreeNode node in selectedNodes)
                                {
                                    if (!nodesToSelect.Contains(node))
                                        nodesToUnselect.Add(node);
                                    else
                                        nodesToSelect.Remove(node);
                                }

                                foreach (var node in nodesToSelect)
                                    node.Selected = true;
                                foreach (var node in nodesToUnselect)
                                    node.Selected = false;
                            }
                        }
                        else
                            unSelectAll = true;
                    }
                }
                if (unSelectAll)
                {
                    UnselectAll();
                    SelectableTreeNode node = (e.Node as SelectableTreeNode);
                    if (node != null)
                    {
                        if (!selectedNodes.Contains(node))
                            selectedNodes.Add(node);
                        node.Selected = true;
                    }
                }
                base.OnAfterSelect(e);
            }
            finally
            {
                if ((Control.ModifierKeys & Keys.Control) != 0 || (!preventMultiSelect && (Control.ModifierKeys & Keys.Shift) != 0))
                    EndUpdate();
            }
        }
        protected override void OnEnter(EventArgs e)
        {
            lastNode = null;
            base.OnEnter(e);
        }
        public TreeNode GetLastLeafNode(TreeNode node)
        {
            if (node == null)
                return null;
            if (node.Nodes.Count == 0)
                return node;
            return GetLastLeafNode(node.Nodes[node.Nodes.Count - 1]);
        }
        public TreeNode GetPreviousNode(TreeNode current, bool wrap = true)
        {
            if (current == null)
                return null;

            if (current.PrevNode != null)
                return GetLastLeafNode(current.PrevNode);
            if (current.Parent == null)
                return wrap ? GetLastLeafNode(Nodes[Nodes.Count - 1]) : null;
            return current.Parent;
        }
        public TreeNode GetNextNode(TreeNode current, bool lookInChilds, bool wrap = true)
        {
            if (current == null)
                return null;

            if (lookInChilds && current.Nodes.Count != 0)
                return current.Nodes[0];
            if (current.NextNode != null)
                return current.NextNode;
            if (current.Parent == null)
                return wrap ? Nodes[0] : null;

            return GetNextNode(current.Parent, false);
        }


        /// <summary>
        /// Returns true if node1 is after node2
        /// </summary>
        /// <param name="node1"></param>
        /// <param name="node2"></param>
        /// <returns>true if node1 is after node2, false otherwise</returns>
        public bool NodeIsAfter(TreeNode node1, TreeNode node2)
        {
            if (node1 == null || node2 == null)
                return false;
            int startingLevel1 = node1.Level;
            int startingLevel2 = node2.Level;
            while (node1.Level > node2.Level)
                node1 = node1.Parent;
            while (node2.Level > node1.Level)
                node2 = node2.Parent;
            while (node1.Parent != node2.Parent)
            {
                node1 = node1.Parent;
                node2 = node2.Parent;
            }
            return node1.Index == node2.Index ? startingLevel1 > startingLevel2 : node1.Index > node2.Index;

        }
        public SelectableTreeNode[] GetNodesBetweenNodes(SelectableTreeNode start, SelectableTreeNode end, bool visibleOnly)
        {
            List<SelectableTreeNode> nodes = new List<SelectableTreeNode>();
            if (start == null || end == null)
                return nodes.ToArray();
            if (start == end)
            {
                nodes.Add(start);
                return nodes.ToArray();
            }
            if (NodeIsAfter(start, end))
            {
                var aux = end;
                end = start;
                start = aux;
            }

            if (visibleOnly && !(end.Parent == null || end.Parent.IsExpanded))
                return nodes.ToArray();
            nodes.Add(start);
            SelectableTreeNode current = start;
            do
            {
                current = GetNextNode(current, !visibleOnly || current.IsExpanded, false) as SelectableTreeNode;
                if (current != null)
                    nodes.Add(current);
            }
            while (current != null && current != end);

            return nodes.ToArray();
        }

        protected override void OnBeforeCollapse(TreeViewCancelEventArgs e)
        {
            BeginUpdate();
            base.OnBeforeCollapse(e);
        }
        protected override void OnAfterCollapse(TreeViewEventArgs e)
        {

            (e.Node as SelectableTreeNode).CollapseNode();
            base.OnAfterCollapse(e);
            EndUpdate();
        }

        protected override void OnBeforeExpand(TreeViewCancelEventArgs e)
        {
            BeginUpdate();
            base.OnBeforeExpand(e);
        }
        protected override void OnAfterExpand(TreeViewEventArgs e)
        {
            
            SelectableTreeNode node = e.Node as SelectableTreeNode;
            if (node != null)
                node.OnExpand();
            if (e.Node.Nodes.Count == 1 && !e.Node.Nodes[0].IsExpanded)
                e.Node.Nodes[0].Expand();

            base.OnAfterExpand(e);
            EndUpdate();
        }
    }
}
